﻿/******************************************************************************
 * Project: Processor Affinity Manager                                        *
 * File:    Shortcut.cs                                                       *
 * Author:  Aaron Jones                                                       *
 * Version: 1                                                                 *
 *                                                                            *
 * Description:                                                               *
 *    This file defines the shortcut program used to execute a process at a   *
 *    given processor affinity.                                               *
 ******************************************************************************/

using System;
using System.Diagnostics;
using System.Windows.Forms;
using System.ComponentModel;
using System.ServiceProcess;

namespace ProcessorAffinityManager
{
    /// <summary>
    /// Allows a process to be started with a specific processor affinity
    /// via a Windows shortcut.
    /// </summary>
    public class Shortcut
    {
        /// <summary>
        /// Exeported executable name so other modules can know how to reference
        /// the shortcut.
        /// </summary>
        public static readonly string ExecutableName =
            System.Reflection.Assembly.GetExecutingAssembly().GetName().Name + ".exe";

        [STAThread]
        static void Main(string[] commandLine)
        {
            // Display the usage if no parameters were given.
            if (commandLine.Length == 0)
            {
                MessageBox.Show(
                    "This shortcut is not meant to be executed directly.",
                    "Invalid Usage",
                    MessageBoxButtons.OK);
            }

            // There are currently 2 combinations of parameters,
            // 1. <executable> <parameters>
            // 2. <affinity> <executable> <parameters>
            // The only way to distinguish between the two cases is test if
            // the first parameter is an integer.  If it is, then case 2.
            else
            {
                Process myProcess = new Process();
                int processors;

                // Passed in a processor affinity.
                try
                {
                    processors = int.Parse(commandLine[0]);
                    myProcess.StartInfo.FileName = commandLine[1];
                    myProcess.StartInfo.Arguments = "";
                    for (int i = 2; i < commandLine.Length; i++)
                        myProcess.StartInfo.Arguments += commandLine[i] + " ";
                }

                // Didn't pass in a processor affinity.  Bring up the picker dialog.
                catch
                {
                    Picker picker = new Picker();
                    if (picker.ShowDialog() == DialogResult.OK)
                    {
                        processors = picker.ProcessorAffinity;
                        myProcess.StartInfo.FileName = commandLine[0];
                        myProcess.StartInfo.Arguments = "";
                        for (int i = 1; i < commandLine.Length; i++)
                            myProcess.StartInfo.Arguments += commandLine[i] + " ";
                    }

                    else
                        return;
                }

                try
                {
                    int procCount = Environment.ProcessorCount;

                    // Make sure the processors we've selected are available.
                    if (processors == 0)
                        throw new Exception("Zero was given.  A processor must be selected");

                    // Make sure the user hasn't selected any processors that
                    // don't exist.
                    for (int i = procCount + 1; i < sizeof(int) * 8; i++)
                        if ((processors & (1 << i - 1)) != 0)
                            throw new Exception("Processor " + i.ToString() + " is not available");

                    // Tell the service to ignore the process we're about to start.
                    try
                    {
                        ServiceController controller = new ServiceController(Service.Name);
                        controller.ExecuteCommand(Service.IgnoreNextProcessCommand);
                    }

                    catch
                    {
                    }

                    try
                    {
                        myProcess.Start();
                        myProcess.ProcessorAffinity = (IntPtr)(processors);
                    }

                    catch(Exception e)
                    {
                        MessageBox.Show(
                            "There was a proccesor affinity error: " + e.Message + ".  The application did not start, or the affinity was not set.",
                            "Processor Affinity Error",
                            MessageBoxButtons.OK);
                    }
                }

                catch (Exception e)
                {
                    MessageBox.Show(
                        "There was a proccesor affinity error: " + e.Message + ".  The application will not start.",
                        "Processor Affinity Error",
                        MessageBoxButtons.OK);
                }
            }
        }
    }
}
